import numpy as np
import functools

class Problem(object):
    class State(object):
        pass

    def init_states(self):
        '''获取初始状态集'''
        pass


    def next_states(self, state):
        '''获取当前状态的后续状态'''
        pass

    def next_actions(self, state):
        '''获取当前状态能够可以采取的步骤'''
        pass

    def is_goal_state(self, state):
        '''判断状态是否为目标状态'''
        pass

    def move(self, state, action):
        '''当前状态采取动作得到新状态'''
        pass

    pass

class Puzzle(Problem):
    class State(object):
        def __init__(self, size, data):
            self.size = size
            self.data = data
            assert(Puzzle.State.check(size, data))
        
        @staticmethod   
        def check(size, data):
            if size*size != len(data): return False

            for e in data:
                if not (0 <= e and e < size*size): return False


            inversion = 0
            for i in range(len(data)):
                for j in range(i+1, len(data)):
                    if data[i] == data[j]: return False
                    if data[i] > data[j]:
                        inversion += 1
            inversion += data.index(0) // size +  data.index(0) % size
            if inversion % 2 != (size+1)%2: return False

            return True

        def equal(self, state):
            return self.data == state.data

        def __str__(self):
            return str(self.data)
        def __repr__(self):
            return self.__str__()

        def __eq__(self, s):
            return type(s) == Puzzle.State and self.data == s.data

        def __hash__(self):
            import functools
            return functools.reduce(lambda x,y: self.size*self.size*x + y, self.data)

        def move(self, direct):

            zero_pos = self.data.index(0)
            
            d = None
            if direct == 'right':
                if zero_pos % self.size != 0:
                    d = self.data.copy()
                    d[zero_pos] = d[zero_pos-1]
                    d[zero_pos-1] = 0
            elif direct == 'left':
                if zero_pos % self.size != self.size-1:
                    d = self.data.copy()
                    d[zero_pos] = d[zero_pos+1]
                    d[zero_pos+1] = 0

            elif direct == 'down':
                if zero_pos >= self.size:
                    d = self.data.copy()
                    d[zero_pos] = d[zero_pos-self.size]
                    d[zero_pos-self.size] = 0
            elif direct == 'up':
                if zero_pos < self.size*self.size - self.size:
                    d = self.data.copy()
                    d[zero_pos] = d[zero_pos+self.size]
                    d[zero_pos+self.size] = 0

            return Puzzle.State(self.size, d) if d != None else None


        @staticmethod
        def random(size):

            while True:
                result = []

                pools = [i for i in range(size*size)]

                for i in range(size*size):
                    import random
                    p = random.randrange(0, len(pools))
                    result.append(pools[p])
                    del pools[p]
                
                if Puzzle.State.check(size, result):
                    return Puzzle.State(size, result)
            return None
            


    def __init__(self, size, init=None):
        self.size = size
        self.init_state = init if init != None else Puzzle.State.random(size)
        self.goal_state = Puzzle.State(size, [i for i in range(1, size*size)] + [0])

    def init_states(self):
        return [self.init_state]


    def next_states(self, state):

        return [ e for e in [state.move(d) for d in ['left', 'right', 'up', 'down']] if e != None ]

    def is_goal_state(self, state):
        return self.goal_state == state

(XY0, XY1, YZ0, YZ1, ZX0, ZX1) = (0, 1, 2, 3, 4, 5)
class MagicCube(Problem):
    class State(Problem.State):
        def __init__(self, size, data):
            assert(MagicCube.State.check(size, data))
            (self.size, self.data) = (size, data)
            

        def move(self, action):
            import copy
            assert(action[0] in ['x','y','z'] and action[1] in range(self.size))
            d = copy.deepcopy(self.data)
            if action[0] == 'x':
                MagicCube.State.rotate_x(self.size, d, action[1])
            elif action[0] == 'y':
                MagicCube.State.rotate_y(self.size, d, action[1])
            else:   # action[0] = 'z'
                assert(action[0] == 'z')
                MagicCube.State.rotate_z(self.size, d, action[1])

            return MagicCube.State(self.size, d)

        def __eq__(self, s):
            return type(s) == MagicCube.State and len(self.data) == len(s.data) and \
                   functools.reduce(lambda x,y: x and y, [(self.data[i] == s.data[i]).all() for i in range(len(s.data))], True)

        def __hash__(self):
            return hash(functools.reduce(lambda x,y: x+y, [e for d in self.data for r in d for e in r]))

        def __repr__(self):
            return functools.reduce(lambda x,y: x+y, [e for d in self.data for r in d for e in r]) 

        def print(self, side):
            if side == XY0:
                for i in range(self.size):
                    print(self.data[XY0][i,:])
            elif side == XY1:
                for i in range(self.size):
                    print(self.data[XY1][self.size-1-i,:])
            elif side == YZ0:
                for i in range(self.size):
                    print(self.data[YZ0][self.size-1-i,range(self.size-1,-1,-1)])
            elif side == YZ1:
                for i in range(self.size):
                    print(self.data[YZ1][self.size-1-i,:])
            elif side == ZX0:
                for i in range(self.size):
                    print(self.data[ZX0][:,self.size-1-i])
            elif side == ZX1:
                for i in range(self.size):
                    print(self.data[ZX1][:,i])


        @staticmethod
        def rotate_x(s, d, i):
            MagicCube.State.rotate_side(s, [d[XY0], d[ZX1], d[XY1], d[ZX0]], i)
            if i == 0: MagicCube.State.rotate_surface(s, d[YZ0])
            if i == s-1: MagicCube.State.rotate_surface(s, d[YZ1])

        @staticmethod
        def rotate_y(s, d, i):
            MagicCube.State.rotate_side(s, [d[YZ0], d[XY1], d[YZ1], d[XY0]], i)

            if i == 0: MagicCube.State.rotate_surface(s, d[ZX0])
            if i == s-1: MagicCube.State.rotate_surface(s, d[ZX1])

        @staticmethod
        def rotate_z(s, d, i):
            MagicCube.State.rotate_side(s, [d[ZX0], d[YZ1], d[ZX1], d[YZ0]], i)
            if i == 0: MagicCube.State.rotate_surface(s, d[XY0])
            if i == s-1: MagicCube.State.rotate_surface(s, d[XY1])
            

        @staticmethod
        def rotate_side(s, d, i):
            t = d[0][:,i].copy()
            d[0][:,i] = d[3][i,range(s-1,-1,-1)]
            d[3][i,:] = d[2][:,i]
            d[2][:,i] = d[1][i,range(s-1,-1,-1)]
            d[1][i,:] = t
    
        @staticmethod
        def rotate_surface(s, d):
            t = np.array(d)
            for i in range(s):
                d[:,s-1-i] = t[i,:]


        @staticmethod
        def check(size, data):
            return True

        
        @staticmethod
        def goal(size):
            '''xy(0,1), yz(2,3) zx(4,5), '''

            #return MagicCube.State(size, [np.ones([size, size])*i for i in range(6)])
            return MagicCube.State(size, [np.array([[i]*size]*size) for i in ['w','b','y','r','o','g']])

        @staticmethod
        def check(size, data):
            return True

    def __init__(self, size, init=None):
        self.size = size
        self.init_state = init if init != None else Puzzle.State.random(size)
        self.goal_state = MagicCube.State.goal(size)

    def next_states(self, state):
        return [state.move(a) for a in [[r,i] for r in ('x','y','z') for i in range(self.size)]]

    def init_states(self):
        return [self.init_state]

    def is_goal_state(self, state):
        return self.goal_state == state


def breadth_first_search(prob, max_breadth=5000000):
    from collections import deque

    state_queue = deque([(None,None)])
    visited_states = set()

    while 0 < len(state_queue) and len(state_queue) < max_breadth:
        p = state_queue.popleft()

        s = prob.next_states(p[-1]) if p[-1] != None else prob.init_states()
        s = [ i for i in s if i not in visited_states]

        if s != []:
            for e in s :
                if prob.is_goal_state(e):
                    result = [e]
                    while p[-1] != None:
                        result.append(p[-1])
                        p = p[0]
                    result.reverse()
                    return result
                
                else:
                    state_queue.append((p, e))

    return None



def depth_first_search(prob, max_depths=5000000):
    state_stack = [prob.init_states()]

    visited_states = set(state_stack[-1])
    
    print('init: '+ str(prob.init_state))

    while 0 < len(state_stack) and len(state_stack) < max_depths:
        
        if len(state_stack[-1]) == 0:
            state_stack.pop()
        else:
            if prob.is_goal_state(state_stack[-1][-1]):
                return [s[-1] for s in state_stack]
            else:
                s = prob.next_states(state_stack[-1][-1])
                s = [ i for i in s if i not in visited_states]

                if s != []:
                    state_stack.append(s)
                    visited_states.update(s)
                else:
                    state_stack[-1].pop()
    return None

if __name__ == '__main__':

    #path = breadth_first_search(Puzzle(3, Puzzle.State(3,[1,2,3,4,5,7,8,6,0])))
    

    g = MagicCube.State.goal(2)
    s0 = MagicCube.State(2, [
        np.array([['g','b'],['o','y']]), 
        np.array([['b','y'],['w','o']]), 
        np.array([['r','w'],['y','r']]), 
        np.array([['w','r'],['o','g']]),
        np.array([['y','o'],['r','g']]),
        np.array([['b','g'],['b','w']])])

    g = g.move(['z',1])
    g = g.move(['x',1])
    #s0 = s0.move(['z', 1])
    #s0.print(XY0)
    #print()
    #s0.print(XY1)
    #print()
    #s0.print(YZ0)
    #print()
    #s0.print(YZ1)
    #print()
    #s0.print(ZX0)
    #print()
    #s0.print(ZX1)
    
    print(s0 == s0.move(['z',0]).move(['z',0]).move(['z',0]).move(['z',0]))
    
    prob = MagicCube(2, g)

    path = depth_first_search(prob)
    print(path)
